import cv2
import numpy as np
from pprint import pprint as pp
from tkinter import filedialog as fd
import os

# 修改记录：
#      1. 2024-08-09： 调整二值化检测阈值，由于新的图片加入了透明度，导致文字灰度增加，暂时将透明度的阈值
#                      从`30`调整为 `70`
#      2. 2024-08-31:  检测点改为黑色方块，修改蓝色区块识别阈值（后续可能仍然会继续修改）
#      3. 2024-09-07:  修改蓝色标签检测参数
#      4. 2024-09-09:  修改检测正方形比例的参数，将长宽比的比值从 `5` 改为 `4.7`
#      5. 2024-10-08:  增加二值图像预处理方法，调整方向识别权重，调整轮廓面积阈值为 `40000`
#      6. 2024-10-09:  调整正反检测算法：计算图片两端的权重，权重大的为左边，权重小的为右边
#      7. 2024-10-11:  调整检测宽高比，改为 `3`，增加调试开关
# TODO:
#      1. 暂时无法处理图像反光情况，处理方向：从亮度通道换成其他颜色通道
class ImageProcessor:
    
    # 初始化
    def __init__(self):
        
        self.image_path = None
        self.origin_image = None
        self.tmp_image = None
        self.hsv_image = None
        self.target_image = None
        self.result_image = None
        self.number_image = None # 读数图片
        self.code_image = None # 编码图片
        self.target_contour = None  # 目标轮廓
        self.target_rect = None
        self.image_angle = -1
        self.window_w = 500
        self.debug = False
        # self.lower_blue = [41, 90, 90]
        # self.upper_blue = [147, 255, 255]
        self.lower_blue = [35, 43, 46]
        self.upper_blue = [77, 255, 255]
        # self.main_win = window

    def process_x(self):
        self.process(self.image_path)

    # 处理指定图片
    def process(self, path):
        
        self.loadImage(path)

        contour = self.detectBlueRegion(self.hsv_image)
        angle, _ = self.calcAngleAndRect(contour)
        result_image = self.correction(self.origin_image, angle)

        if self.debug:
            cv2.imshow("result tmp:", result_image)

        tmp_img = self.hsvImg(result_image)
        tmp_contour = self.detectBlueRegion(tmp_img)
        _, rect = self.calcAngleAndRect(tmp_contour)

        # box = cv2.boxPoints(rect).reshape((-1, 1, 2)).astype(np.int32)
        # cv2.polylines(result_image, [box], True, (0, 255,255), 10)

        pt, size, _ = rect
        print("================ pt", pt, size)
        size = np.array(size).astype(np.int32)
        w, h = size[0], size[1]
        if w < h:
            w, h = h, w

        sub_img = cv2.getRectSubPix(result_image, (w, h), pt)
        self.number_image = sub_img

        # 设置最终处理图片
        pos = 1
        if self.needFlip(sub_img, (w, h)):
            # print("do flip")
            self.result_image = self.roateImg(result_image, 180)
            self.target_image = self.roateImg(sub_img, 180)
        else:
            self.target_image = sub_img
            self.result_image = result_image

        code_img = cv2.getRectSubPix(result_image, (int(w*0.7), int(h*0.9)), (pt[0], pt[1]-pos*int(h)-30))
        if pos == 1:
            self.code_image = code_img
        else:
            self.code_image = self.roateImg(code_img, 180)

        if self.debug:
            cv2.imshow("code_img", self.code_image)

    # 加载图片
    def loadImage(self, path):
        self.image_path = path
        
        # 检查文件是否存在
        if not os.path.exists(path):
            print(f"错误: 找不到文件 {path}")
            raise FileNotFoundError(f"找不到文件: {path}")
            
        # 加载图片并检查是否成功
        image = cv2.imread(path)
        if image is None:
            print(f"错误: 无法加载图片 {path}")
            raise ValueError(f"无法读取图片文件: {path}")
            
        self.origin_image = image
        self.tmp_image = image.copy()
        self.hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    
    # 检测蓝色区域
    def detectBlueRegion(self, hsv_image):
        lower_blue = np.array(self.lower_blue)
        upper_blue = np.array(self.upper_blue)

        # 创建掩码
        mask = cv2.inRange(hsv_image, lower_blue, upper_blue)

        # 增加平滑运算，减少图像噪点和毛刺
        opt = 7
        kernel = np.ones((opt, opt), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.GaussianBlur(mask, (opt, opt), 0)

        # win_mask = cv2.imshow("mask_win", cv2.resize(mask, np.array([300, 300])))
        # cv2.resizeWindow('mask_win', 300, 300)

        # 找到目标轮廓
        max_area = 0
        max_id = -1
        contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        for i, contour in enumerate(contours):
            area = cv2.contourArea(contour)
            if area < 40000:
                continue

            rect = cv2.minAreaRect(contour)
            (cx, cy), (l, w), theta = rect
            # print("@@@@@@@@@@@2", area, l, w, theta)
            if l < w:
                l, w = w, l

            # 在临时图像上绘制轮廓
            if self.debug:
                cv2.drawContours(self.tmp_image, contours, i, (0, 255, 0), 5)
                # box = cv2.boxPoints(rect)
                # box = np.int0(box)      
                # cv2.drawContours(self.tmp_image, [box], 0, (0,0,255), 2)
                self.showImg("tmp", self.tmp_image)
            
            delta = l / w 

            if self.debug:
                print("delta:", delta)

            if delta < 3:
                continue

            if area > max_area:
                max_area = area
                max_id = i

        # 符合条件的轮廓
        if self.debug:
            print("max id:", max_id)

        return contours[max_id]

    # 计算倾斜角度
    def calcAngleAndRect(self, contour):
        rect = cv2.minAreaRect(contour)
        (cx, cy), (l, w), theta = rect

        cv2.drawContours(self.tmp_image, contour, -1, (0, 255, 255), 3)
        # 初步计算水平放置角度
        flag = 1
        if l < w:
            theta = 90-theta
            flag = -1

        return flag*theta, rect
    
    # 校正图片方向
    def correction(self, image, angle):
        return self.roateImg(image, angle)

    # 获取 HSV 图片
    def hsvImg(self, image):
        return cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 计算像素值
    def pix_sum(self, col_data):
        white = 0
        black = 0

        for v in col_data:
            if v==0:
                black = black+1
            else:
                white = white+1
                
        return [white, black]

    def countBlack(self, gray_img, h):
        black = 0
        for i in range(h):
            col = gray_img[:, i]
            for v in col:
                if v == 0:
                    black = black + 1
        return black
    
    # 判断箭头的位置
    def checkArrow(self, data, h):
        gate = 0.45
        score = 0
        for i in data:
            if i !='-' and i > gate:
                score = score + 1

        print("score ", score, " h: ", h, "  h*0.45  ", h*gate)
        return score < h*gate
 
    # 判断是否需要翻转
    def needFlip(self, image, size):
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        v_channel = hsv_image[:, :, 2]  # 亮度通道
        _, gray = cv2.threshold(v_channel, 120, 255, cv2.THRESH_BINARY)
        
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(7, 7))
        gray = cv2.erode(gray, kernel)

        kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(7, 7))
        gray = cv2.dilate(gray, kernel)

        if self.debug:
            cv2.imshow("gray", gray) 

        w, h = size
        if w < h: 
            w, h = h, w
        
        opt1 = 11
        th = int(h*3/4)
        tmp_gray = cv2.getRectSubPix(gray, (th, th), (h/2, h/2))
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(opt1, opt1))
        tmp_gray = cv2.dilate(tmp_gray, kernel)
        if self.debug:
            cv2.imshow("tmp_gray", tmp_gray)

        tmp_gray1 = cv2.getRectSubPix(gray, (th, th), (w-h/2, h/2))
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(opt1, opt1))
        tmp_gray1 = cv2.dilate(tmp_gray1, kernel)
        if self.debug:
            cv2.imshow("tmp_gray1", tmp_gray1)

        # kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(7, 7))
        # gray = cv2.dilate(gray, kernel)

        black_left = self.countBlack(tmp_gray, th)
        black_right = self.countBlack(tmp_gray1, th)

        return black_left < black_right

    # 旋转图形
    def roateImg(self, img, angle):
        (h, w) = img.shape[:2]

        # 定义旋转中心点
        center = (w // 2, h // 2)

        # 获取旋转矩阵
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)

        # 计算旋转后的图像尺寸
        cos = np.abs(rotation_matrix[0, 0])
        sin = np.abs(rotation_matrix[0, 1])

        # 计算新图像的宽度和高度
        new_w = int((h * sin) + (w * cos))
        new_h = int((h * cos) + (w * sin))

        # 调整旋转矩阵以考虑平移
        rotation_matrix[0, 2] += (new_w / 2) - center[0]
        rotation_matrix[1, 2] += (new_h / 2) - center[1]

        # 应用旋转矩阵
        return cv2.warpAffine(img, rotation_matrix, (new_w, new_h))
    
    def showImg(self, title, img):
        h, w = img.shape[:2]

        rate = h/w*1.0
        size = np.array((self.window_w , self.window_w*rate)).astype(np.int32)

        resized_img = cv2.resize(img, size)
        cv2.imshow(title, resized_img)
        cv2.resizeWindow(title, size[0], size[1])

    # 显示原始图片
    def showOriginImage(self):
        self.showImg("main_win", self.origin_image)

    # 显示 HSV 图片
    def showHsvImage(self):
        self.showImg("HSV Image", self.hsv_image)

    def showResult(self):
        self.showImg("Result Image", self.result_image)

    def showTarget(self):
        self.showImg("Target Image", self.target_image)

    def wait(self):
        while True:
            key = cv2.waitKey(50)
            if key == ord('q'):
                cv2.destroyAllWindows()
                break
            if key == ord('f'):
                fname = fd.askopenfilename()
                print('file name:', fname)